"""
往图像上添加水印

"""


import os
import time
import numpy as np
import pickle
import struct
import re
import cv2
from tqdm import tqdm
import imageio
import json
from PIL import Image
from collections import defaultdict
import random
import copy
from ctypes import CDLL,byref,create_string_buffer,cdll
from PIL import Image, ImageColor, ImageFont, ImageDraw, ImageFilter
import sys
from multiprocessing import Process

from tool.data_enhance import DataEnhance
from tool import utils, imgaug_tool, filesystem, opencv_tool, via_tool # export PYTHONPATH=$PYTHONPATH:`pwd` 




class WaterMark:
    def __init__(self,mask_dir):

        self.load_background_image( mask_dir)
        print("load_background_image ok...")


    def load_background_image(self,mask_dir):

        maskp = filesystem.get_all_filepath(mask_dir, [".jpg", ".png"])
        self.mask_cv_image = [cv2.imread(f) for f in maskp]

    def random_get_mask_image(self, img_shape, fix_background=False):

        cv_image = copy.copy(self.mask_cv_image[np.random.randint(len(self.mask_cv_image))])

        if cv_image.shape[1] > img_shape[1] and cv_image.shape[0] > img_shape[0]:
            start_width = np.random.randint(cv_image.shape[1] - img_shape[1])
            start_height = np.random.randint(cv_image.shape[0] - img_shape[0])
        else:
            cv_image = cv2.resize(cv_image, (img_shape[1], img_shape[0]))
            start_height = 0
            start_width = 0
        new_image = cv_image[start_height:start_height+ img_shape[0], start_width:start_width+img_shape[1] ]
        return new_image

    def run(self, all_files, dynamic_weight_by_filesize=False):
        

        for f in tqdm(all_files, total=len(all_files)):
            if np.random.randint(4) == 0:continue

            image = cv2.imread(f, 1)
            start = np.random.randint(image.shape[1] // 3)
            end = np.random.randint(image.shape[1] // 2, image.shape[1])
            mask = self.random_get_mask_image(( image.shape[0], end - start))
            # mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)

            if dynamic_weight_by_filesize:
                file_size_ratio = (os.path.getsize(f) - 8000) / (18000 - 8000)
                ratio = np.clip(file_size_ratio, 0, 0.4)
            else:
                ratio = np.random.randint(2,65) / 100
            crop_img = image[:,start:end ]
            new_img = crop_img *(1-ratio) + mask * ratio
            new_img = np.array(new_img, dtype=np.uint8)
            image[:,start:end ] = new_img

            cv2.imwrite(f, image)


def marker_for_driving_license(data_dir, mask_dir, thread_count,
                                    dw_bs):
    mask_adder = WaterMark(mask_dir)

    all_files = filesystem.get_all_filepath(data_dir, [".jpg"])
    length = len(all_files)
    per_thread_num = length // thread_count
    
    thread_list = []
    for i in range(thread_count):
        t = Process(target=mask_adder.run, args=(all_files[i*per_thread_num: (i+1)*per_thread_num], dw_bs, ))
        t.start()


if __name__ == "__main__":

    thread_count = 12
    mask_dir = "/mnt/disk1/vanlance/project/driving_license/bg/mask"
    data_dir = "/home/swls/work_dir/ocr/code/syn/driving_license_374k"
    dynamic_weight_by_filesize=True
    marker_for_driving_license(data_dir, mask_dir, thread_count, 
        dynamic_weight_by_filesize)



    # data_ori_path = r"/home/swls/work_dir/ocr/code/yolo_train/yolo_driving_license/gen/180000/via_region_data_ori.json"
    # via_path =      r"/home/swls/work_dir/ocr/code/yolo_train/yolo_driving_license/gen/180000/via_region_data_size.json"
    # via_tool.convert_to_via(data_ori_path, via_path)


